from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool, StructuredTool, tool


class MultiplyCalculatorInput(BaseModel):
    a: str = Field(description="the first number.")
    b: str = Field(description="the second number.")


def multiply(a: str, b: str) -> int:
    """multiply two numbers."""
    return len(a) * len(b)


multiply_calculator = StructuredTool.from_function(
    func=multiply,  # 工具具体逻辑
    name="multiply_Calculator",  # 工具名
    description="multiply two numbers",  # 工具信息
    args_schema=MultiplyCalculatorInput,  # 工具接受参数信息
    return_direct=True,  # 直接作为工具的输出返回给调用者
    handle_tool_error=True,  # 报错了继续执行，不会吧那些报错行抛出，也可以自定义函数处理，handle_tool_error=函数名
)

# print(calculator.invoke("计算2和3的乘积"))
# print(multiply_calculator.invoke({"a": 3, "b": 2}))
# print(multiply_calculator.invoke("3 * 2"))